from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

class RAGService:
    def __init__(self, model_name='all-mpnet-base-v2'):
        self.model = SentenceTransformer(model_name)
        self.index = None
        self.documents = []

    def embed_documents(self, documents):
        self.documents = documents
        embeddings = self.model.encode(documents, convert_to_numpy=True)
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(embeddings)

    def search(self, query, top_k=5):
        query_embedding = self.model.encode([query], convert_to_numpy=True)
        distances, indices = self.index.search(query_embedding, top_k)
        results = [(self.documents[i], distances[0][j]) for j, i in enumerate(indices[0])]
        return results